#!/usr/bin/env python
#
#  rrd.py
#  Simple RRDTool wrapper
#  Copyright (c) 2008 Corey Goldberg (corey@goldb.org)
#
#  Download the Windows version of RRDTool from:
#    http://www.gknw.net/mirror/rrdtool/
# 
#  You may need these fonts if RRDTool throws an error when you graph:
#    http://dejavu.sourceforge.net/wiki/index.php/Main_Page


import os
import time
import random


class RRD(object):
    def __init__(self, rrd_name):
        self.rrd_name = rrd_name

    def create_rrd(self, interval, graph_config):
		interval_mins = float(interval) / 60  
		heartbeat = int(interval) * 2
		ds_string = ''
		
		datasets = self.get_datasets(graph_config)
		
		args = self.config_value('create_args', graph_config)
		if args is None: args = ''
		
		for ds_name in datasets:
			ds_type = self.ds_value(ds_name, 'type', graph_config)
			if ds_type is None: ds_type = 'GAUGE'
			
			ds_min = self.ds_value(ds_name, 'min', graph_config)
			if ds_min is None: ds_min = 'U'
			
			ds_max = self.ds_value(ds_name, 'max', graph_config)
			if ds_max is None: ds_max = 'U'
			
			# continue with constr.
			ds_string += ' DS:%s:%s:%d:%s:%s' % (ds_name, ds_type, heartbeat, ds_min, ds_max)
		cmd_create = ''.join((
			'rrdtool create ', self.rrd_name, ' --step ', str(interval), ds_string,
			' RRA:AVERAGE:0.5:1:', str(int(4000 / interval_mins)),
			' RRA:AVERAGE:0.5:', str(int(30 / interval_mins)), ':800',
			' RRA:AVERAGE:0.5:', str(int(120 / interval_mins)), ':800',
			' RRA:AVERAGE:0.5:', str(int(1440 / interval_mins)), ':800',
			' ' + args
			))
		#print cmd_create
		cmd = os.popen4(cmd_create)
		cmd_output = cmd[1].read()
		for fd in cmd: fd.close()
		if len(cmd_output) > 0:
			raise RRDException, 'Unable to create RRD: ' + cmd_output
    
    def update(self, values, floats=True):   
    	vs = []
	for v in values:
	  if floats:
	   vs.append(str(float(v)))
	  else:
	   vs.append(str(int(v)))

        values_args = ':'.join(vs)
        cmd_update = 'rrdtool update %s N:%s' % (self.rrd_name, values_args)
	cmd = os.popen4(cmd_update)
        cmd_output = cmd[1].read()
        for fd in cmd: fd.close()
        if len(cmd_output) > 0:
	    if floats:
	      return self.update(values, False)
            raise RRDException, 'Unable to update RRD: ' + cmd_output
    
    def random_color(self, range1=100, range2=200):
        res = '#'
	random.seed()
        while len(res) < 7:
	 res += hex(random.randrange(range1, range2))[2:][:1]
	return res

    def rgb(self, c):
        split = (c[0:2], c[2:4], c[4:6])
        return [int(x, 16) for x in split]

    def lumdiff(self, colA, colB):
        r1, g1, b1 = self.rgb(colA[1:]);
        r2, g2, b2 = self.rgb(colB[1:]);
        
        l1 = 0.2126 * pow(float(r1)/255.0, 2.2) + 0.7152 * pow(float(g1)/255.0, 2.2) + 0.0722 * pow(float(b1)/255, 2.2);
        l2 = 0.2126 * pow(float(r2)/255.0, 2.2) + 0.7152 * pow(float(g2)/255.0, 2.2) + 0.0722 * pow(float(b2)/255, 2.2);
        
        if l1 > l2:
            return (l1 + 0.05) / (l2+0.05)
        else:
            return (l2 + 0.05) / (l1+0.05)
    
    def getminlum(self, colours, col):
        min = 100;
        if len(colours) == 0:
          return 5
        for c in colours:
           l = self.lumdiff(col, c)
           if l < min:
               min = l
        return min
 
    def colordelta(self, col, delta):
	r,g,b = self.rgb(col[1:])
	r+=delta
	g+=delta
	b+=delta
	if r>255: r=255
	if g>255: g=255
	if b>255: b=255
	if r<0: r=0
	if g<0: g=0
	if b<0: b=0
	res = '#%02x%02x%02x'%(r,g,b)
	return res

    def config_value(self, key, config):
        for c in config:
	  if type(c) == list and c[0] == key:
	    return c[1]
	return None
     
    def ds_value(self, ds, key, config):
        for c in config:
	  if type(c) == list and len(c) >=3 and c[0] == ds and c[1] == key:
	    return c[2]
	return None
    
    def get_datasets(self, config):
        ds = []
	for c in config:
	  if type(c) == list and not c[0] in ['graph_title', 'create_args', 'graph_args', 'graph_category', 'graph_info', 'graph_order', 'graph_vlabel', 'graph_total', 'graph_scale', 'graph', 'host_name', 'update', 'graph_period', 'graph_vtitle', 'service_order', 'graph_width', 'graph_height', 'graph_printfformat']:
            if not c[0] in ds:
	     ds.append(c[0])
	return ds

    def graph(self, mins, graph_config, output_filename=None, title2=None):       
		start_time = 'now-%s' % (mins * 60)
		if output_filename is None:
			output_filename = self.rrd_name + '.png'
		end_time = 'now'
		width = '500'
		height = '150'

		
		title = self.config_value('graph_title', graph_config)
		if title is None:
		  title = title2
		else:
		  title = title2 + ' / ' + title
		
		args = self.config_value('graph_args', graph_config)
		if args is None: args = ''
		
		vlabel = self.config_value('graph_vlabel', graph_config)
		if vlabel is None: vlabel = 'value'

		
		cur_date = time.strftime('%m/%d/%Y %H\:%M\:%S', time.localtime())
		defs = ''
		
		datasets = self.get_datasets(graph_config)
		colours = '#00CC00 #0066B3 #FF8000 #FFCC00 #330099 #990099 #CCFF00 #FF0000 #808080 #008F00 #00487D #B35A00 #B38F00 #6B006B #8FB300 #B30000 #BEBEBE #80FF80 #80C9FF #FFC080 #FFE680 #AA80FF #EE00CC #FF8080 #666600 #FFBFFF #00FFCC #CC6699 #999900'.split()
		cur_col = 0
		
		for ds_name in datasets:
			rfn = os.path.commonprefix([os.path.abspath(os.curdir), self.rrd_name])
			# relpath workaround
			#print "XXXXX", [os.path.abspath(os.curdir), self.rrd_name], rfn
			#rfn = self.rrd_name[len(rfn)+1:]
			rfn = os.path.basename(self.rrd_name)
			
			ds_colour = self.ds_value(ds_name, 'colour', graph_config)
			# find good contrast colours
                        if ds_colour is None:
                          ds_colour = colours[cur_col]
			  cur_col += 1
			  if cur_col >= len(colours): cur_col -= len(colours)
                        ds_colour2 = self.colordelta(ds_colour, 100)
		
			ds_label = self.ds_value(ds_name, 'label', graph_config)
			if ds_label is None: ds_label = ds_name
			
			ds_draw = self.ds_value(ds_name, 'draw', graph_config)
			if ds_draw is None: ds_draw = 'LINE2'
			
			ds_hrule = self.ds_value(ds_name, 'line', graph_config)
			if ds_hrule is None:
			  ds_hrule = ''
			else:
			  ds_hrule = ' HRULE:%s:%s' % (ds_hrule, ds_colour)
		
			ds_negative = self.ds_value(ds_name, 'negative', graph_config)
			ds_negative = (not ds_negative is None)
			
			ds_cdef = self.ds_value(ds_name, 'cdef', graph_config)
		        if ds_cdef is None: ds_cdef = ''


			ds_disable = self.ds_value(ds_name, 'graph', graph_config)
			ds_disable = (not ds_disable is None)

			# we use mods, take care!
			# ouch, stuff got complex!
			ds_name_add = ''
                        defs += ' DEF:' + ds_name + '=' + rfn + ':' + ds_name + ':AVERAGE'
			defs += ' DEF:' + ds_name + '_avg=' + rfn + ':' + ds_name + ':AVERAGE:step=%d'%(mins*5)

			if ds_cdef != '':
			   dscmod = ds_cdef.split(',')
			   dscmod[0] = dscmod[0] + ds_name_add
			   dscmod = ','.join(dscmod)
			   defs += ' CDEF:' + ds_name + ds_name_add + '_mod=' + dscmod

			   dscmod = ds_cdef.split(',')
			   dscmod[0] = dscmod[0] + '_avg' + ds_name_add
			   dscmod = ','.join(dscmod)
			   defs += ' CDEF:' + ds_name + '_avg' + ds_name_add + '_mod=' + dscmod
			   
			   ds_name_add += '_mod'

			if ds_negative:
			   defs += ' CDEF:' + ds_name + ds_name_add +'_neg=' + ds_name + ds_name_add + ',-1,*'
			   defs += ' CDEF:' + ds_name + '_avg' + ds_name_add + '_neg=' + ds_name + '_avg' + ds_name_add + ',-1,*'
			   ds_name_used = ds_name + '_neg'


			if True:
			   # only draw the average line if its a line!
			   if ds_draw in ['LINE', 'LINE1', 'LINE2', 'LINE3']:
			     defs +=    ' LINE2:' + ds_name + '_avg' + ds_name_add + ds_colour2 

			   defs +=    ' ' + ds_draw + ':' + ds_name + ds_name_add + ds_colour + ':' + '"%-20s'%ds_label + '"'
			   defs += 	' VDEF:' + ds_name + ds_name_add + 'last=' + ds_name + ds_name_add + ',LAST' + \
					' VDEF:' + ds_name + ds_name_add + 'avg=' + ds_name + ds_name_add + ',AVERAGE' + \
					' VDEF:' + ds_name + ds_name_add + 'max=' + ds_name + ds_name_add + ',MAXIMUM' + \
					' VDEF:' + ds_name + ds_name_add + 'min=' + ds_name + ds_name_add + ',MINIMUM' + \
					' GPRINT:' + ds_name + ds_name_add + 'last:"%10.2lf%S"' + \
					' GPRINT:' + ds_name + ds_name_add + 'max:"%10.2lf%S"' + \
					' GPRINT:' + ds_name + ds_name_add + 'avg:"%10.2lf%S"' + \
					' GPRINT:' + ds_name + ds_name_add + 'min:"%10.2lf%S\l"' + \
					' ' + ds_hrule
		
		cmd_graph = 'rrdtool graph ' + output_filename + \
			' COMMENT:"                     "' + \
			' COMMENT:"       Last"' + \
			' COMMENT:"    Maximum"' + \
			' COMMENT:"    Average"' + \
			' COMMENT:"        Min\l"' + \
			' ' + defs + \
			' COMMENT:"' + cur_date + '"' + \
			' -W "CREATED BY MONITORDATASINK"' + \
			' --title="' + title +'"' + \
			' --vertical-label="' + vlabel + '"' \
			' --start=' + start_time + \
			' --end=' + end_time + \
			' --width=' + width + \
			' --height=' + height + \
			' --slope-mode --lazy ' + args
			# RE-ADD --lazy
		#print '>>>>>>>>>>>>>>>>', rfn
		#print cmd_graph
  		dir = os.getcwd()
		os.chdir(os.path.dirname(self.rrd_name))
		cmd = os.popen4(cmd_graph)
		os.chdir(dir)
        	cmd_output = cmd[1].read()
		if len(cmd_output) > 10:
		  print "local error:"
		  print 'output: ', cmd_output
		  print 'cmd: ', cmd_graph
                  print 'config: ', graph_config
		for fd in cmd: fd.close()
            
            
class RRDException(Exception): pass
